package vn.edu.hcmut.cse.ai.search.problems.blocksworld;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import vn.edu.hcmut.cse.ai.search.local.HeuristicFunction;
import vn.edu.hcmut.cse.ai.search.problems.blocksworld.predicate.BlocksWorldPredicate;
import vn.edu.hcmut.cse.ai.search.problems.blocksworld.predicate.On;
import vn.edu.hcmut.cse.ai.search.problems.blocksworld.predicate.OnTable;

public class BlocksWorldSecondHeuristicFunction implements HeuristicFunction {

    @Override
    public double heuristicEvaluation(Object currentState, Object goalState) {
        BlocksWorldState blocksWorldCurrentState = (BlocksWorldState) currentState;
        BlocksWorldState blocksWorldGoalState = (BlocksWorldState) goalState;

        // For each block that is on a correct substructure, add 1 point for 
        // every block in the substructure, and for each block that is on an 
        // incorrect substructure, subtract 1 point for every block in the 
        // substructure
        HashMap<Integer, Integer> onMap = blocksWorldCurrentState.getOnMap();
        List<BlocksWorldPredicate> substructureList = new ArrayList<>();
        int heuristicPoint = 0;
        for (Integer onBlock : onMap.keySet()) {
            substructureList.clear();

            // Search for a substructure below the onBlock
            boolean isSubstructureCorrect = true;
            Integer underBlock;
            
            // Add the (onBlock, underBlock) to the substructure
            underBlock = onMap.get(onBlock);            
            substructureList.add(new On(onBlock, underBlock));

            // While loop for searching the substructure
            while (true) {
                // Check if underBlock is on another block or on table
                onBlock = underBlock;
                
                if (onMap.containsKey(onBlock)) {
                    underBlock = onMap.get(onBlock);
                    On onPred = new On(onBlock, underBlock);
                    substructureList.add(onPred);
                } else {
                    // underBlock is on table
                    OnTable onTablePred = new OnTable(onBlock);
                    substructureList.add(onTablePred);
                    break;
                }
            }

            // Check if the substructure is correct
            for (BlocksWorldPredicate pred : substructureList) {
                if (blocksWorldGoalState.getPredicates().contains(pred) == false) {
                    isSubstructureCorrect = false;
                }
            }

            // Mark the current substructure
            if (isSubstructureCorrect == true) {
                heuristicPoint = heuristicPoint + (substructureList.size() - 1);
            } else {
                heuristicPoint = heuristicPoint - (substructureList.size() - 1);
            }
        }

        return (double) heuristicPoint;
    }
}
